import torch
import torch.nn as nn
import torch.nn.functional as F

class AttentionLayer(nn.Module):
    """
    注意力机制层，用于让模型关注序列中的重要部分
    """
    def __init__(self, hidden_dim):
        super(AttentionLayer, self).__init__()
        self.attention = nn.Linear(hidden_dim, 1)
        
    def forward(self, gru_output):
        """
        :param gru_output: GRU的输出，形状为 (batch_size, seq_len, hidden_dim)
        :return: 加权后的上下文向量和注意力权重
        """
        # 计算注意力分数
        attention_scores = self.attention(gru_output).squeeze(-1)  # (batch_size, seq_len)
        
        # 应用softmax得到注意力权重
        attention_weights = F.softmax(attention_scores, dim=1)  # (batch_size, seq_len)
        
        # 使用注意力权重计算上下文向量
        context_vector = torch.bmm(
            attention_weights.unsqueeze(1),  # (batch_size, 1, seq_len)
            gru_output  # (batch_size, seq_len, hidden_dim)
        ).squeeze(1)  # (batch_size, hidden_dim)
        
        return context_vector, attention_weights

class AttentiveGRUScoringModel(nn.Module):
    """
    带有注意力机制的GRU评分模型，用于处理多轮对话的特征向量序列。
    模型结构：GRU层 -> 注意力层 -> 全连接回归头 -> Sigmoid激活 + 缩放
    """
    def __init__(self, input_dim=384, hidden_dim=128, n_layers=2, dropout=0.2):
        """
        模型初始化
        :param input_dim: 输入特征向量的维度
        :param hidden_dim: GRU 隐藏层的维度
        :param n_layers: GRU 层的数量
        :param dropout: Dropout 比率，用于防止过拟合
        """
        super(AttentiveGRUScoringModel, self).__init__()
        
        self.gru = nn.GRU(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=n_layers,
            batch_first=True,
            dropout=dropout if n_layers > 1 else 0,
            bidirectional=True  # 使用双向GRU以捕获更多上下文信息
        )
        
        # 双向GRU，隐藏状态维度翻倍
        self.attention = AttentionLayer(hidden_dim * 2)
        
        self.regressor = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 1)
        )
        
    def forward(self, x, lengths):
        """
        模型的前向传播
        :param x: 经过填充（padded）的输入张量，维度为 (batch_size, max_seq_len, input_dim)
        :param lengths: 一个包含批次中每个序列原始长度的张量，用于处理填充
        :return: 一个0-100之间的评分值，维度为 (batch_size, 1)
        """
        # 为了让GRU忽略填充部分，我们首先打包序列
        packed_input = nn.utils.rnn.pack_padded_sequence(x, lengths.cpu(), batch_first=True, enforce_sorted=False)
        
        # GRU的输出是打包后的序列
        packed_output, _ = self.gru(packed_input)
        
        # 解包序列，得到完整的输出序列和有效长度
        gru_output, _ = nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True)
        
        # 应用注意力机制
        context_vector, attention_weights = self.attention(gru_output)
        
        # 将上下文向量输入回归头
        score = self.regressor(context_vector)
        
        # 压缩到0-100
        final_score = torch.sigmoid(score) * 100
        
        return final_score, attention_weights

# 为了保持向后兼容性，我们保留原来的类名，但内部使用增强的模型
class ComprehensiveScoringModel(nn.Module):
    """
    一个综合评分模型，用于处理多轮对话的特征向量序列。
    现在内部使用带有注意力机制的GRU模型。
    """
    def __init__(self, input_dim=384, hidden_dim=128, n_layers=2, dropout=0.2):
        super(ComprehensiveScoringModel, self).__init__()
        self.attentive_model = AttentiveGRUScoringModel(input_dim, hidden_dim, n_layers, dropout)
        
    def forward(self, x, lengths):
        """
        模型的前向传播
        :param x: 经过填充（padded）的输入张量，维度为 (batch_size, max_seq_len, input_dim)
        :param lengths: 一个包含批次中每个序列原始长度的张量，用于处理填充
        :return: 一个0-100之间的评分值，维度为 (batch_size, 1)
        """
        score, _ = self.attentive_model(x, lengths)
        return score 